/*
 *  Disk.cpp
 *
 *  CDisk class for direct disk operatoin
 *
 *  Author: Binghua Duan <bhduan@gmail.com>
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License version 2 as
 *  published by the Free Software Foundation.
 */

#include "stdafx.h"
#include "Disk.h"

CDisk::CDisk(LPCTSTR lpszPhyDrv)
{
	m_hPhyDrv = NULL;
	m_bReadOnly = FALSE;
	Open(lpszPhyDrv);
}

CDisk::~CDisk()
{
	Close();
}

BOOL CDisk::Open(LPCTSTR lpszPhyDrv)
{
	enum { SHIFT_GB = 30 };

	Close();
	m_hPhyDrv = CreateFile(lpszPhyDrv, GENERIC_ALL, FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
	if (m_hPhyDrv != INVALID_HANDLE_VALUE) {
		DISK_GEOMETRY pdg;
		if (GetGeometry(&pdg)) {
			m_nTotalSectors = pdg.Cylinders.QuadPart * pdg.TracksPerCylinder * pdg.SectorsPerTrack;
			m_nBytesPerSector = pdg.BytesPerSector;
			m_strPhyDrv = lpszPhyDrv;

			if (m_strPhyDrv == _T("\\\\.\\PHYSICALDRIVE0") ||
				((ULONGLONG)m_nTotalSectors * m_nBytesPerSector >> SHIFT_GB) >= 40) {
					/* NOTE: DISABLE THE WRITING TO HARD-DISK!!! */
					m_bReadOnly = TRUE;
			}

			return TRUE;
		} else {
			Close();
		}
	}

	m_hPhyDrv = NULL;
	m_strPhyDrv.Empty();
	return FALSE;
}

void CDisk::Close()
{
	if (m_hPhyDrv) {
		CloseHandle(m_hPhyDrv);
		m_hPhyDrv = NULL;
		m_strPhyDrv.Empty();
	}
}

BOOL CDisk::GetGeometry(DISK_GEOMETRY *pdg) const
{
	if (!IsValid())
		return FALSE;

	DWORD junk; // discard results
	BOOL bRet = DeviceIoControl(m_hPhyDrv, IOCTL_DISK_GET_DRIVE_GEOMETRY, NULL, 0, pdg, sizeof(*pdg),
		&junk, (LPOVERLAPPED) NULL);

	return bRet;
}

BOOL CDisk::Write(LPCVOID lpBuffer, UINT nNumberOfBytesToWrite)
{
	if (!IsValid() || m_bReadOnly)
		return FALSE;

	DWORD nNumberOfBytesWritten = 0;
	BOOL bRet = WriteFile(m_hPhyDrv, lpBuffer, nNumberOfBytesToWrite, &nNumberOfBytesWritten, 0);
	ASSERT(bRet || nNumberOfBytesWritten != nNumberOfBytesToWrite);
	return nNumberOfBytesWritten == nNumberOfBytesToWrite;
}

BOOL CDisk::WriteSector(LPCVOID lpBuffer, UINT nNumberOfSectorsToWrite)
{
	return Write(lpBuffer, nNumberOfSectorsToWrite * m_nBytesPerSector);
}

BOOL CDisk::Read(LPVOID lpBuffer, UINT nNumberOfBytesToRead)
{
	if (!IsValid())
		return FALSE;

	DWORD nNumberOfBytesRead = 0;
	BOOL bRet = ReadFile(m_hPhyDrv, lpBuffer, nNumberOfBytesToRead, &nNumberOfBytesRead, 0);
	ASSERT(bRet || nNumberOfBytesToRead != nNumberOfBytesRead);
	return nNumberOfBytesToRead == nNumberOfBytesRead;
}

BOOL CDisk::ReadSector(LPVOID lpBuffer, UINT nNumberOfSectorsToRead)
{
	return Read(lpBuffer, nNumberOfSectorsToRead * m_nBytesPerSector);
}

BOOL CDisk::Seek(ULONGLONG nSeekPos)
{
	if (!IsValid())
		return FALSE;

	LONG dwHighPos = (LONG)(nSeekPos >> 32);
	DWORD dwLowPos = SetFilePointer(m_hPhyDrv, (DWORD)nSeekPos, &dwHighPos, FILE_BEGIN);
	return dwLowPos != INVALID_SET_FILE_POINTER || GetLastError() != NO_ERROR;
}

BOOL CDisk::SeekSector(UINT nSector)
{
	return Seek((ULONGLONG)nSector * m_nBytesPerSector);
}

BOOL CDisk::Eject()
{
	PREVENT_MEDIA_REMOVAL pmr;
	DWORD dwNotUsed;

	pmr.PreventMediaRemoval = FALSE;

	BOOL b1 = DeviceIoControl(m_hPhyDrv, IOCTL_STORAGE_MEDIA_REMOVAL, &pmr, sizeof(pmr), NULL, 0, &dwNotUsed, NULL);
	BOOL b2 = DeviceIoControl(m_hPhyDrv, IOCTL_STORAGE_EJECT_MEDIA, NULL, 0, NULL, 0, &dwNotUsed, NULL);
	return b1 && b2;
}
